# coding=utf-8
# Chern-C
# Copyright (C) 2024 NJU-SAML
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import os
import re
from typing import Dict, Iterable, List, Mapping, Optional, Set, Tuple

from chernc.core.cargo import CargoConfig
from chernc.slicer.code_chunk import CodeChunk, CodeChunkID
from chernc.slicer.code_slice import CodeSlice

Marks = Dict[CodeSlice, Tuple[str, str, bool]]


def text_of_marks(marks: Marks) -> str:
    slices_text = ""
    for s in sorted(marks.keys(), key=lambda s: s.level):
        slices_text += marks[s][0] + "\n"
        slices_text += s.text
        slices_text += marks[s][1] + "\n\n"
    return slices_text


class MarkedCodeChunk(CodeChunk):
    def __init__(self, chunk: CodeChunk) -> None:
        self.id: str = chunk.id
        self.slices: List[CodeSlice] = chunk.slices
        self.size: int = chunk.size
        self.w_depends: Dict[CodeChunkID, int] = chunk.w_depends
        self.w_depended: Dict[CodeChunkID, int] = chunk.w_depended
        self.level = chunk.level
        self.marks: Marks = {s: (self.getStartMark(s), self.getEndMark(s), self.isOptional(s)) for s in chunk.slices}
        self.optionals = {s for s in chunk.slices if self.isOptional(s)}

    @staticmethod
    def getStartMark(slice: CodeSlice):
        return f"// CHERNC_MARK: {os.path.abspath(slice.file)}:{slice.start_line}"

    @staticmethod
    def getEndMark(slice: CodeSlice):
        return f"// CHERNC_END_MARK: {os.path.abspath(slice.file)}:{slice.end_line}"

    @staticmethod
    def from_codechunk(chunk: CodeChunk):
        new_chunk = MarkedCodeChunk(chunk)
        return new_chunk

    @staticmethod
    def isOptional(slice: CodeSlice):
        return slice.isTypedef()

    @property
    def text(self) -> str:
        return text_of_marks(self.marks)


class SnippetResult(object):
    """
    A snippet represents a result for some slices, these slices cannot be seperated because of how LLM translates them.
    A translation result of a chunk is not always possble to be split in terms of slices
    For example, some slices has relevant semantics and LLM translates them into one unsplitable code
    In this case, those slices (and moreover, the files they belong to) must always be in the same place
    """

    def __init__(self, slices: Iterable[CodeSlice], text: str, names: Iterable[str]):
        self.slices = list(slices)
        self.text = text
        self.names = list(names)

    def add_pub(self, name):
        kinds = [
            "fn",
            "struct",
            "enum",
            "trait",
            "const",
            "static",
            "type",
        ]

        for kind in kinds:
            self.text = re.sub(
                r'(\b(?:unsafe|safe)\s+(?:extern\s+"cdecl"|extern\s+"C"|extern)?)?\b' + kind + r"\s+" + name + r"\b",
                r"pub \g<0>",
                self.text,
                re.M,
            )

        self.text = re.sub(r"\bpub\s+pub\b", "pub", self.text, re.M)


class UseManager:
    class Mod:
        children: Dict[str, "UseManager.Mod"]
        used: bool
        name: str
        as_name: str

        def __init__(self, name: str):
            self.name = name
            self.children = {}
            self.used = False
            self.as_name = name

        def add(self, name) -> "UseManager.Mod":
            if name not in self.children:
                self.children[name] = UseManager.Mod(name)
            return self.children[name]

        def use(self, as_name: str):
            self.used = True
            self.as_name = as_name

        def dont_use(self):
            self.used = False

        def paths(self) -> List[str]:
            paths = []
            used_together = [child for child in self.children.values() if child.used and child.as_name == child.name]
            used_as = [child for child in self.children.values() if child.used and child.as_name != child.name]
            if len(used_together) == 1:
                paths.append(f"{used_together[0].name}")
            elif len(used_together) > 1:
                paths.append(f"{{{', '.join([used.name for used in used_together])}}}")

            for used in used_as:
                paths.append(f"{used.name} as {used.as_name}")

            for child in self.children.values():
                paths += child.paths()
            return [f"{self.name}::{path}" for path in paths]

        def copy(self) -> "UseManager.Mod":
            mod = UseManager.Mod(self.name)
            mod.used = self.used
            mod.as_name = self.as_name
            for child in self.children.values():
                mod.children[child.name] = child.copy()
            return mod

        def __or__(self, other: "UseManager.Mod") -> "UseManager.Mod":
            if self.name != other.name:
                raise Exception(f"Cannot merge two mods with different names: {self.name} and {other.name}")
            mod = UseManager.Mod(self.name)
            for child in self.children.keys() | other.children.keys():
                if child in self.children and child in other.children:
                    mod.children[child] = self.children[child] | other.children[child]
                elif child in self.children:
                    mod.children[child] = self.children[child].copy()
                else:
                    mod.children[child] = other.children[child].copy()
            mod.used = self.used or other.used
            if other.as_name != self.name:
                mod.as_name = other.as_name
            else:
                mod.as_name = self.as_name
            return mod

        def __sub__(self, other: "UseManager.Mod") -> "UseManager.Mod":
            if self.name != other.name:
                return self.copy()
            mod = UseManager.Mod(self.name)
            mod.used = self.used and not other.used
            mod.as_name = self.as_name
            for child in self.children.keys():
                if child not in other.children:
                    mod.children[child] = self.children[child].copy()
                else:
                    mod.children[child] = self.children[child] - other.children[child]
            return mod

    root: Mod
    useds: Dict[str, Mod]

    def __init__(self):
        self.root = UseManager.Mod("")
        self.useds = {}

    @staticmethod
    def parse_use_statement(mod_path) -> Tuple[List[str], List[str], Optional[str]]:
        """
        for a::b::{c, d} get ["a", "b"] and ["c", "d"]
        """
        used_pattern = r"((?P<used>\w+)|\{(?P<useds>[\w\s,]+)\})"
        path_pattern = r"(?P<path>(\w+::)*)"
        as_pattern = r"(\s+as\s+(?P<as>\w+))?"
        paths = []
        useds = []
        as_name = None
        if m := re.match(path_pattern + used_pattern + as_pattern, mod_path):
            paths: List[str] = m["path"].split("::")[:-1]
            useds = m["used"] or m["useds"]
            useds = list(map(str.strip, useds.split(",")))
            as_name = m["as"]
        return paths, useds, as_name

    def use(self, mod_path: str):
        paths, useds, as_name = self.parse_use_statement(mod_path)
        cur = self.root
        for mod in paths:
            cur = cur.add(mod)
        if "self" in useds:
            useds.remove("self")
            self.use("::".join(paths))
        if not useds:
            return
        if as_name and as_name in self.useds:
            return

        for used in useds:
            if not as_name and used in self.useds:
                return
            added = cur.add(used)
            name = as_name or used
            added.use(name)
            self.useds[name] = added

    def dont_use(self, name: str):
        if name in self.useds:
            self.useds[name].dont_use()
            del self.useds[name]

    def use_declarations(self) -> str:
        declarations = [f"use {path};" for path in self.paths()]
        return "\n".join(declarations)

    def paths(self):
        return [path[2:] for path in self.root.paths()]

    def __or__(self, other: "UseManager") -> "UseManager":
        manager = UseManager()
        manager.root = self.root | other.root

        manager._walk(manager.root, lambda mod: manager.useds.setdefault(mod.as_name, mod) if mod.used else None)
        return manager

    def __sub__(self, other: "UseManager") -> "UseManager":
        manager = UseManager()
        manager.root = self.root - other.root
        manager._walk(manager.root, lambda mod: manager.useds.setdefault(mod.as_name, mod) if mod.used else None)
        return manager

    def copy(self) -> "UseManager":
        manager = UseManager()
        manager.root = self.root.copy()
        manager._walk(manager.root, lambda mod: manager.useds.setdefault(mod.as_name, mod) if mod.used else None)
        return manager

    def _walk(self, mod: "UseManager.Mod", fn):
        fn(mod)
        for child in mod.children.values():
            self._walk(child, fn)


class ModDeclaration(object):
    def __init__(self, mod_name: str, modifier: str = "pub"):
        self.mod_name = mod_name
        self.modifier = modifier

    def __str__(self):
        if len(self.modifier) == 0:
            return f"mod {self.mod_name};"
        else:
            return f"{self.modifier} mod {self.mod_name};"

    def __repr__(self):
        return str(self)

    def __eq__(self, value: object) -> bool:
        if not isinstance(value, ModDeclaration):
            return False
        return self.mod_name == value.mod_name

    def __hash__(self):
        return hash(self.mod_name) + hash(self.modifier)


class ExternDeclaration(object):
    def __init__(self, crate_name: str):
        self.crate_name = crate_name

    def __str__(self):
        return f"extern crate {self.crate_name};"

    def __repr__(self):
        return str(self)

    def __eq__(self, value: object) -> bool:
        if not isinstance(value, ExternDeclaration):
            return False
        return self.crate_name == value.crate_name

    def __hash__(self):
        return hash(self.crate_name)


class FileResult(object):
    def __init__(self, path: str):
        self.path = path
        self.mod_declarations: List[ModDeclaration] = list()
        self.extern_decls: List[ExternDeclaration] = list()
        self.uses: UseManager = UseManager()
        self.snippets: List[SnippetResult] = []
        self.isEntry = False

    def add_snippet(self, snippet: SnippetResult):
        self.snippets.append(snippet)

    def move(self, path: str):
        self.path = path

    @property
    def text(self) -> str:
        out = ""
        # extern crates declarations
        for decl in self.extern_decls:
            out += str(decl) + "\n"
        # Mod names
        for mod_decl in self.mod_declarations:
            out += str(mod_decl) + "\n"
        out += "\n"

        # Use decls
        out += self.uses.use_declarations() + "\n"

        # Code
        out += "\n".join(s.text for s in self.snippets) + "\n"
        return out

    def get_text_with_details(self, mark_position: bool, mark_source: bool) -> str:
        out = ""
        # extern crates declarations
        for decl in self.extern_decls:
            out += str(decl) + "\n"
        # Mod names
        for mod_decl in self.mod_declarations:
            out += str(mod_decl) + "\n"
        out += "\n"

        # Use decls
        out += self.uses.use_declarations() + "\n"

        # Code
        for snippet in self.snippets:
            for s in snippet.slices:
                if mark_position:
                    out += MarkedCodeChunk.getStartMark(s) + "\n"
                if mark_source:
                    out += "\n".join(f"// {line}" for line in s.text.splitlines(keepends=False)) + "\n"
            out += snippet.text + "\n"
            for s in snippet.slices:
                if mark_position:
                    out += MarkedCodeChunk.getEndMark(s) + "\n"
        return out + "\n"


class DummySnippetResult(SnippetResult):
    """
    This is a dummy result for a single slice.
    """

    def __init__(self, slice: CodeSlice):
        super().__init__([slice], f"// Dummy translation for {slice.loc}", [decl[1] for decl in slice.decls])


Match = Mapping[CodeSlice, SnippetResult]


class ChunkResult:
    """
    This represents a result for a chunk.
    """

    chunk: CodeChunk
    result_table: Dict[CodeSlice, SnippetResult]
    results: Set[SnippetResult]

    def __init__(self, chunk: CodeChunk, result_table: Match):
        self.chunk = chunk
        self.result_table = dict(result_table)
        for slice in chunk.slices:
            if slice not in result_table:
                self.result_table[slice] = DummySnippetResult(slice)
        self.results = set(self.result_table.values())

    @property
    def text(self):
        return "\n".join(r.text for r in self.results)


class SnippetRepairResult(SnippetResult):
    def __init__(self, slices: Iterable[CodeSlice], text: str, names: Iterable[str]):
        super().__init__(slices, text, names)


class FileRepairResult(FileResult):
    def __init__(self, path: str):
        super().__init__(path)
        # self.repaired_snippets: List[SnippetRepairResult] = []

        # 一些内部对象到行数range映射，开始行-结束行
        self.extern_decls_range: Dict[ExternDeclaration, Tuple[int, int]] = {}
        self.mod_decls_range: Dict[ModDeclaration, Tuple[int, int]] = {}
        self.use_range: Tuple[int, int] = (0, 0)
        self.snippets_range: Dict[SnippetResult, Tuple[int, int]] = {}

        self._setup_line_mappings()

    def _setup_line_mappings(self):
        line_num = 0
        # extern crates declarations
        for decl in self.extern_decls:
            extern_decl_text = str(decl) + "\n"
            self.extern_decls_range[decl] = (line_num, line_num + extern_decl_text.count("\n"))
            line_num += extern_decl_text.count("\n")
        # Mod names
        for mod_decl in self.mod_declarations:
            mod_decl_text = str(mod_decl) + "\n"
            self.mod_decls_range[mod_decl] = (line_num, line_num + mod_decl_text.count("\n"))
            line_num += mod_decl_text.count("\n")

        newline_text = "\n"
        line_num += newline_text.count("\n")

        # Use decls
        use_text = self.uses.use_declarations() + "\n"
        self.use_range = (line_num, line_num + use_text.count("\n"))
        line_num += use_text.count("\n")
        # Code
        for s in self.snippets:
            snippet_text = s.text + "\n"
            self.snippets_range[s] = (line_num, line_num + snippet_text.count("\n"))
            line_num += snippet_text.count("\n")

    def find_snippet_by_line(self, line: int) -> Optional[SnippetResult]:
        for snippet, (start, end) in self.snippets_range.items():
            if line >= start and line < end:
                return snippet
        return None

    @staticmethod
    def from_file_result(file_result: FileResult) -> "FileRepairResult":
        repair_result = FileRepairResult(file_result.path)
        repair_result.mod_declarations = file_result.mod_declarations.copy()
        repair_result.extern_decls = file_result.extern_decls.copy()
        repair_result.uses = file_result.uses.copy()
        repair_result.snippets = file_result.snippets.copy()
        repair_result.isEntry = file_result.isEntry
        repair_result._setup_line_mappings()
        return repair_result


class ProjectRepairResult:
    """
    This represents a repair result for a project.
    """

    cargo: CargoConfig
    files: Dict[str, FileRepairResult]

    def __init__(self, cargo: CargoConfig, files: Dict[str, FileRepairResult]):
        self.cargo = cargo
        self.files = files


class ProjectResult:
    """
    This represents a result for a project.
    """

    cargo: CargoConfig
    files: Dict[str, FileResult]

    def __init__(self, cargo: CargoConfig, files: Dict[str, FileResult]):
        self.cargo = cargo
        self.files = files

    def to_project_repair_result(self) -> ProjectRepairResult:
        return ProjectRepairResult(self.cargo, {k: FileRepairResult.from_file_result(v) for k, v in self.files.items()})
